package com.eight.cloud.client.handler;

import com.eight.cloud.client.service.subscribe.SubscribeObject;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.messaging.simp.stomp.*;
import org.springframework.web.socket.WebSocketHttpHeaders;
import org.springframework.web.socket.messaging.WebSocketStompClient;

import java.lang.reflect.Type;
import java.util.Map;
import java.util.concurrent.*;

/**
 * @ClassName: SessionConnectHandler
 * @Author: TXC
 * @Date: 2024-11-19 17:49
 **/
@Slf4j
public class SessionConnectHandler extends StompSessionHandlerAdapter {
    @Setter
    private WebSocketStompClient stompClient;
    @Setter
    private WebSocketHttpHeaders headers;
    @Setter
    private String serverWsUrl;
    private final Map<String, SubscribeObject> subscribeObjects;
    private ScheduledFuture<?> heartbeatFuture;
    public SessionConnectHandler(Map<String, SubscribeObject> subscribeObjects) {
        this.subscribeObjects = subscribeObjects;
    }

    @Override
    public void afterConnected(StompSession session, StompHeaders connectedHeaders) {
        log.info("Connected to server");
        subscribeObjects.forEach((key, subscribeObject) ->
                session.subscribe(subscribeObject.getTopic(), subscribeObject.getSubscribeHandler())
        );
        startHeartbeatMonitor(session);
    }

    @Override
    public void handleFrame(StompHeaders headers, Object payload) {
        log.info("Received message: {}", payload);
    }

    @Override
    public Type getPayloadType(StompHeaders headers) {
        return String.class;
    }

    @Override
    public void handleException(StompSession session, StompCommand command, StompHeaders headers, byte[] payload,
                                Throwable exception) {
        super.handleException(session, command, headers, payload, exception);
        log.error("Error occurred: {}", exception.getMessage());
    }

    @Override
    public void handleTransportError(StompSession session, @NotNull Throwable exception) {
        super.handleTransportError(session, exception);
        log.error("Transport error occurred: {}", exception.getMessage());
        try {
            Thread.sleep(10000);
            // 尝试重新连接
            reconnect(session);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }

    private void startHeartbeatMonitor(StompSession session) {
        if (heartbeatFuture != null && !heartbeatFuture.isCancelled()) {
            // 取消之前的任务
            heartbeatFuture.cancel(true);
        }
        if (stompClient.getTaskScheduler() != null) {
            heartbeatFuture = stompClient.getTaskScheduler().scheduleAtFixedRate(() -> {
                try {
                    // 检查心跳是否正常
                    if (!session.isConnected()) {
                        log.error("Session {} is not connected. Attempting to reconnect...", session.getSessionId());
                        // 重新连接逻辑
                        reconnect(session);
                    } else {
                        log.info("Heartbeat received from server for session {}", session.getSessionId());
                    }
                } catch (Exception e) {
                    log.error("Error during heartbeat monitoring for session {}: {}", session.getSessionId(), e.getMessage());
                }
            }, stompClient.getDefaultHeartbeat()[0]);
        }
    }

    private void reconnect(StompSession session) {
        try {
            if (session != null && session.isConnected()) {
                // 断开连接
                session.disconnect();
                log.info("Disconnected old session: {}", session.getSessionId());
            }
            log.info("Trying to reconnect...");
            StompSession stompSession = stompClient.connect(serverWsUrl, headers, this).get();
            log.info("Reconnected to STOMP server with new session ID: {}", stompSession.getSessionId());
        } catch (Exception e) {
            log.error("Failed to reconnect: {}", e.getMessage());
        }
    }
}
